import os
from pathlib import Path
from transformers import VisionEncoderDecoderModel
from llama_index.core import SimpleDirectoryReader

# 强制所有Hugging Face相关库离线
os.environ["TRANSFORMERS_OFFLINE"] = "1"
# os.environ["HF_DATASETS_OFFLINE"] = "1"  # 如果使用了datasets库
os.environ["LLAMA_INDEX_CACHE_DIR"] = "./cache"  # 设置llama-index的本地缓存

def load_model(model_path: str):
    """加载本地模型"""
    try:
        model = VisionEncoderDecoderModel.from_pretrained(
            model_path,
            local_files_only=True
        )
        print("模型加载成功！")
        return model
    except Exception as e:
        print(f"模型加载失败: {e}")
        return None

def load_documents(directory: str, input_files: list = None):
    """加载文档"""
    try:
        dir_reader = SimpleDirectoryReader(
            input_dir=directory,
            input_files=input_files,
            # 禁用任何可能的网络请求
            required_exts=[".txt", ".pdf", ".docx"]  # 明确指定支持的扩展名
        )
        documents = dir_reader.load_data()
        print(f"文档数量: {len(documents)}")
        if documents:
            print(documents[0].text[:100])  # 打印第一个文档的前100个字符
        return documents
    except Exception as e:
        # 过滤掉网络连接相关的错误
        if "connect" not in str(e).lower() and "http" not in str(e).lower():
            print(f"文档加载失败: {e}")
        return None

# 定义路径（使用Path对象更健壮）
model_path = Path("D:/ideaSpace/MyPython/models/vit-gpt2-image-captioning")
doc_dir = Path("D:/ideaSpace/rag-in-action-master/90-文档-Data/黑悟空")

# 加载模型
model = load_model(model_path)

# 加载目录中的所有文件
documents = load_documents(doc_dir)

# 仅加载特定文件
specific_file = ["D:/ideaSpace/rag-in-action-master/90-文档-Data/黑悟空/设定.txt"]
documents = load_documents(doc_dir, input_files=specific_file)

print("-" * 50)
print("LlamaIndex-构建Document对象")
from llama_index.core import Document

# 创建多个文档对象，并为其添加元数据
documents = [
    Document(
        text="一个充满烈焰和硫磺气息的地下洞窟，火焰从地底不断喷涌，照亮整个深渊。场景中有熔岩河流穿行，燃烧的火山石在空中漂浮。悟空需要利用自己的跳跃能力和金箍棒在熔岩之间穿行，同时对抗来自地狱的火焰妖怪。",
        metadata={
            "filename": "火照深渊场景.md",
            "category": "游戏场景",
            "file_path": "/data/黑悟空/火照深渊场景.md",
            "author": "GameScience",
            "creation_date": "2024-11-20",
            "last_modified_date": "2024-11-21",
            "file_type": "markdown",
            "word_count": 28,
        },
    ),
    Document(
        text="一片高耸入云的山脉，云雾缭绕，风力强劲。悟空需要通过飞跃山崖、利用筋斗云飞行，以及在大风中保持平衡来穿越场景。敌人主要是隐匿在云层中的飞禽妖怪和岩石机关兽。",
        metadata={
            "filename": "风起长空场景.md",
            "category": "游戏场景",
            "file_path": "/data/黑悟空/风起长空场景.md",
            "author": "GameScience",
            "creation_date": "2024-11-20",
            "last_modified_date": "2024-11-21",
            "file_type": "markdown",
            "word_count": 28,
        },
    )]

# 打印每个文档的元数据
for doc in documents:
    print(f"Metadata for {doc.metadata['filename']}:")
    for key, value in doc.metadata.items():
        print(f"  {key}: {value}")
print("-" * 40)

